package org.github.manx98.jlibhac.tools.fssystem;

import cn.hutool.core.util.ByteUtil;
import lombok.extern.slf4j.Slf4j;
import org.github.manx98.jlibhac.common.math.Uint16;
import org.github.manx98.jlibhac.common.math.Uint64;
import org.github.manx98.jlibhac.common.math.UnsignedNumber;
import org.github.manx98.jlibhac.fs.IStorage;

import java.io.IOException;
import java.math.BigInteger;
import java.nio.ByteOrder;

/**
 * @author Manx98
 * @version 1.0
 * @since 2022/10/3 22:07
 */
@Slf4j
public class Aes128CtrStorage extends SectorStorage {
    private static final int BlockSize = 0x10;
    private final Aes128CtrTransform decryptor;
    private final byte[] Counter;
    private final long counterOffset;

    public Aes128CtrStorage(IStorage baseStorage, byte[] key, long counterOffset, byte[] counterHi, boolean leaveOpen) {
        super(baseStorage, BlockSize, leaveOpen);
        if (key == null) throw new IllegalArgumentException("key cannot be null");
        if (key.length != BlockSize) throw new IllegalArgumentException("Key must be " + BlockSize + " bytes long");
        this.counterOffset = counterOffset;
        byte[] initialCounter = new byte[BlockSize];
        if (counterHi != null) {
            System.arraycopy(counterHi, 0, initialCounter, 0, 8);
        }
        this.decryptor = new Aes128CtrTransform(key, initialCounter);
        this.Counter = decryptor.Counter;
    }

    public static byte[] createCounter(Uint64 hiBytes, long offset) {
        byte[] counter = new byte[0x10];
        hiBytes.toByteArray(counter, 0, ByteOrder.BIG_ENDIAN);
        byte[] bytes = ByteUtil.longToBytes(offset / 0x10, ByteOrder.BIG_ENDIAN);
        System.arraycopy(bytes, 0, counter, 8, bytes.length);
        return counter;
    }

    @Override
    public void read(long offset, byte[] destination, int destOffset, int len) throws IOException {
        log.debug("Aes128CtrStorage read " + len + " bytes from " + destOffset);
        super.read(offset, destination, destOffset, len);
        updateCounter(counterOffset + offset);
        decryptor.transformBlock(destination, destOffset, len);
    }

    private void updateCounter(long offset) {
        UnsignedNumber off = new Uint64(offset).shiftRight(4);
        Uint16 magic = new Uint16(0xFF);
        for (int j = 0; j < 0x7; j++) {
            Counter[0x10 - j - 1] = off.and(magic).byteValue();
            off = off.shiftRight(8);
        }

        // Because the value stored in the counter is offset >> 4, the top 4 bits
        // of byte 8 need to have their original value preserved
        Counter[8] = (byte) ((Counter[8] & 0xF0) | off.and(new Uint64(0x0F)).intValue());
    }
}
